import os
import time
import json
import hmac
import hashlib
import base64
import requests
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

class QCloudHunyuan3D:
    def __init__(self):
        self.secret_id = os.getenv('TENCENT_CLOUD_SECRET_ID')
        self.secret_key = os.getenv('TENCENT_CLOUD_SECRET_KEY')
        self.region = os.getenv('TENCENT_CLOUD_REGION', 'ap-guangzhou')
        # 根据官方文档更新API域名
        self.api_endpoint = 'https://ai3d.tencentcloudapi.com'
        
        # 验证必要的配置是否存在
        if not self.secret_id or not self.secret_key:
            raise ValueError("请在.env文件中配置TENCENT_CLOUD_SECRET_ID和TENCENT_CLOUD_SECRET_KEY")
    
    def _generate_signature(self, method, endpoint, params, timestamp, action):
        """
        生成腾讯云API签名
        遵循TC3-HMAC-SHA256算法
        """
        # 1. 生成规范请求字符串
        canonical_uri = endpoint
        canonical_querystring = ''  # 腾讯云3D API使用POST，querystring应为空
        canonical_headers = f"content-type:application/json\nhost:ai3d.tencentcloudapi.com\n"
        signed_headers = "content-type;host"
        
        # 构建请求体
        payload = json.dumps(params)
        hashed_request_payload = hashlib.sha256(payload.encode('utf-8')).hexdigest()
        
        # 构建规范请求
        canonical_request = '\n'.join([
            method,
            canonical_uri,
            canonical_querystring,
            canonical_headers,
            signed_headers,
            hashed_request_payload
        ])
        
        # 2. 生成待签名字符串
        algorithm = "TC3-HMAC-SHA256"
        date = time.strftime('%Y-%m-%d', time.gmtime(int(timestamp)))
        # 更新service名称为ai3d
        credential_scope = f"{date}/ai3d/tc3_request"
        hashed_canonical_request = hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()
        string_to_sign = '\n'.join([
            algorithm,
            timestamp,
            credential_scope,
            hashed_canonical_request
        ])
        
        # 3. 计算签名
        # 计算派生密钥 - 更新service名称为ai3d
        secret_date = hmac.new(f"TC3{self.secret_key}".encode('utf-8'), date.encode('utf-8'), hashlib.sha256).digest()
        secret_service = hmac.new(secret_date, "ai3d".encode('utf-8'), hashlib.sha256).digest()
        secret_signing = hmac.new(secret_service, "tc3_request".encode('utf-8'), hashlib.sha256).digest()
        
        # 计算签名
        signature = hmac.new(secret_signing, string_to_sign.encode('utf-8'), hashlib.sha256).hexdigest()
        
        return signature
    
    def _generate_headers(self, signature, timestamp, action):
        """
        生成请求头
        """
        algorithm = "TC3-HMAC-SHA256"
        date = time.strftime('%Y-%m-%d', time.gmtime(int(timestamp)))
        # 更新service名称为ai3d
        credential_scope = f"{date}/ai3d/tc3_request"
        signed_headers = "content-type;host"
        
        authorization = f"{algorithm} Credential={self.secret_id}/{credential_scope}, SignedHeaders={signed_headers}, Signature={signature}"
        
        return {
            'Content-Type': 'application/json',
            'Host': 'ai3d.tencentcloudapi.com',
            'X-TC-Timestamp': timestamp,
            'X-TC-Version': '2025-05-13',  # 更新API版本
            'X-TC-Region': self.region,
            'X-TC-Action': action,
            'Authorization': authorization
        }
    
    def generate_3d_from_text(self, prompt, options=None):
        """
        根据文本生成3D模型
        :param prompt: 文本描述
        :param options: 可选参数，可包含action字段指定任务类型
        :return: API响应
        """
        try:
            timestamp = str(int(time.time()))
            
            # 从options中获取action参数，如果没有提供则使用默认值
            options = options or {}
            action = options.pop('action', 'SubmitHunyuanTo3DProJob')  # 默认使用专业版
            endpoint = '/'
            
            # 构建请求参数 - 只包含业务参数，不包含公共参数
            params = {
                'Prompt': prompt
            }
            
            # 添加可选参数
            if options:
                params.update(options)
            
            # 生成签名
            signature = self._generate_signature('POST', endpoint, params, timestamp, action)
            
            # 生成请求头
            headers = self._generate_headers(signature, timestamp, action)
            
            # 发送请求
            response = requests.post(
                self.api_endpoint,
                data=json.dumps(params),
                headers=headers
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            return response.json()
            
        except Exception as e:
            print(f"生成3D模型时出错: {str(e)}")
            raise
    
    def generate_3d_from_image(self, image_url, options=None):
        """
        根据图片生成3D模型
        :param image_url: 图片URL或base64编码
        :param options: 可选参数，可包含action字段指定任务类型
        :return: API响应
        """
        try:
            timestamp = str(int(time.time()))
            
            # 从options中获取action参数，如果没有提供则使用默认值
            options = options or {}
            action = options.pop('action', 'SubmitHunyuanTo3DProJob')  # 默认使用专业版
            endpoint = '/'
            
            # 构建请求参数 - 只包含业务参数，不包含公共参数
            params = {
                'ImageUrl': image_url
            }
            
            # 添加可选参数
            if options:
                params.update(options)
            
            # 生成签名
            signature = self._generate_signature('POST', endpoint, params, timestamp, action)
            
            # 生成请求头
            headers = self._generate_headers(signature, timestamp, action)
            
            # 发送请求
            response = requests.post(
                self.api_endpoint,
                data=json.dumps(params),
                headers=headers
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            return response.json()
            
        except Exception as e:
            print(f"根据图片生成3D模型时出错: {str(e)}")
            raise
    
    def get_task_status(self, task_id, options=None):
        """
        查询任务状态
        :param task_id: 任务ID
        :param options: 可选参数，可包含action字段指定查询接口类型
        :return: 任务状态响应
        """
        try:
            timestamp = str(int(time.time()))
            
            # 从options中获取action参数，如果没有提供则根据任务ID推断或使用默认值
            options = options or {}
            action = options.pop('action', 'QueryHunyuanTo3DProJob')  # 默认使用专业版查询接口
            endpoint = '/'
            
            # 构建请求参数 - 只包含业务参数，不包含公共参数
            params = {
                'JobId': task_id
            }
            
            # 生成签名
            signature = self._generate_signature('POST', endpoint, params, timestamp, action)
            
            # 生成请求头
            headers = self._generate_headers(signature, timestamp, action)
            
            # 发送请求
            response = requests.post(
                self.api_endpoint,
                data=json.dumps(params),
                headers=headers
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            return response.json()
            
        except Exception as e:
            print(f"查询任务状态时出错: {str(e)}")
            raise